package cn.tx.doubaodemo.service.impl;

import cn.tx.doubaodemo.service.ChatService;
import cn.tx.doubaodemo.vo.AnswerVO;
import com.volcengine.ark.runtime.model.completion.chat.*;
import io.reactivex.schedulers.Schedulers;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.PostConstruct;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

@Service
public class ChatServiceImpl implements ChatService, DisposableBean {

    private ArkService service;

    @PostConstruct
    public void init() {
        // 从系统环境变量中获取apiKey
        String apiKey = System.getenv("ARK_API_KEY");
        // 连接池，用于支持并发调用
        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();

        // 初始化客户端
        service = ArkService.builder()
                .dispatcher(dispatcher)
                .connectionPool(connectionPool)
                .baseUrl("https://ark.cn-beijing.volces.com/api/v3")    // 调用域名，不填默认访问华北 2 (北京)。支持的域名列表：https://www.volcengine.com/docs/82379/1302013
                .apiKey(apiKey)
                .timeout(Duration.ofSeconds(120))   // 超时时间，默认10分钟
                .connectTimeout(Duration.ofSeconds(20)) // 连接超时时间，默认1分钟
                .build();
    }

    @Override
    public AnswerVO getStandardAnswerDoubao(String input) {
        System.out.println("\n----- standard request -----");
        ChatCompletionRequest chatCompletionRequest = buildDoubaoRequest(input, "你是豆包，是由字节跳动开发的 AI 人工智能助手");

        AnswerVO vo = new AnswerVO();
        List<String> result = new ArrayList<>();
        service.createChatCompletion(chatCompletionRequest)
                .getChoices()
                .forEach(choice -> {
                    String content = choice.getMessage().getContent().toString();
                    System.out.println(content);
                    result.add(content);
                });
        vo.setAnswer(result);
        return vo;
    }

    @Override
    public ResponseBodyEmitter getStreamAnswerDoubao(String input) {
        System.out.println("\n----- streaming request -----");
        ChatCompletionRequest streamChatCompletionRequest = buildDoubaoRequest(input, "你是豆包，是由字节跳动开发的 AI 人工智能助手");

        final ResponseBodyEmitter emitter = new ResponseBodyEmitter(120 * 1000L);
        ExecutorService executorService = Executors.newSingleThreadExecutor();
        executorService.execute(() -> {
            try {
                service.streamChatCompletion(streamChatCompletionRequest)
                        .doOnError(Throwable::printStackTrace)
                        .blockingForEach(
                                choice -> {
                                    if (!choice.getChoices().isEmpty()) {
                                        for (ChatCompletionChoice choiceValue : choice.getChoices()) {
                                            String content = choiceValue.getMessage().getContent().toString();
                                            System.out.print(content);
                                            emitter.send(content, MediaType.TEXT_PLAIN);
                                        }
                                    }
                                }
                        );

            } catch (Exception e) {
                emitter.completeWithError(e);
            }
            emitter.complete();
        });
        return emitter;
    }

    @Override
    public AnswerVO getAnswerAsyncDoubao(String input) {
        // 突然插播的jdk12新特性：switch表达式
        String newInput = switch (input) {
            // 多个条件可以合并，并且不再需要break，因为取消了fall-through
            // 可以作为有返回值的表达式
            case "a", "b" -> "哈哈哈";
            default -> null;
        };
        // 这是jdk13的新特性：yield表达式，作用和上面12的箭头一样
        String newInput1 = switch (input) {
            // 多个条件可以合并，并且不再需要break，因为取消了fall-through
            // 可以作为有返回值的表达式
            case "a":
                yield "哈哈哈";
            case "b":
                yield "嘎嘎嘎";
            default: {
                yield "咔咔咔";
            }
        };

        // 回到正轨继续书写正常的代码
        System.out.println("\n----- streaming request -----");
        ChatCompletionRequest streamChatCompletionRequest = buildDoubaoRequest(input, "你是豆包，是由字节跳动开发的 AI 人工智能助手");

        AnswerVO vo = new AnswerVO();
        List<String> result = new ArrayList<>();
        service.streamChatCompletion(streamChatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .subscribeOn(Schedulers.io())
                .observeOn(Schedulers.single())
                .subscribe(choice -> {
                    if (!choice.getChoices().isEmpty()) {
                        String content = choice.getChoices().getFirst().getMessage().getContent().toString();
                        System.out.println(content);
                        result.add(content);
                    }
                });
        vo.setAnswer(result);
        return vo;
    }

    @Override
    public AnswerVO getDeepSeekAnswer(String input) {
        System.out.println("\n----- standard request from deepSeek -----");
        ChatCompletionRequest chatCompletionRequest = buildDeepSeekRequest(input, "你是deepSeek，你是一个AI编程助手，请帮我输出改正语法后的SQL语句，不要解释直接输出");

        AnswerVO vo = new AnswerVO();
        List<String> result = new ArrayList<>();
        service.createChatCompletion(chatCompletionRequest)
                .getChoices()
                .forEach(choice -> {
                    String content = choice.getMessage().getContent().toString();
                    System.out.println(content);
                    result.add(content);
                });
        vo.setAnswer(result);

        return vo;
    }

    private ChatCompletionRequest buildDoubaoRequest(String input, String... systemInput) {
        return buildRequest("ep-20250219193222-f5n2r", input, systemInput);
    }

    private ChatCompletionRequest buildDeepSeekRequest(String input, String... systemInput) {
        return buildRequest("ep-20250226205422-smdf8", input, systemInput);
    }

    private ChatCompletionRequest buildRequest(String model, String input, String... systemInput) {
        final List<ChatMessage> streamMessages = new ArrayList<>();
        for (String system : systemInput) {
            ChatMessage streamSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(system).build();
            streamMessages.add(streamSystemMessage);
        }
        final ChatMessage streamUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(input).build();
        streamMessages.add(streamUserMessage);

        return ChatCompletionRequest.builder()
                .model(model)
                .messages(streamMessages)
                .build();
    }

    @Override
    public void destroy() throws Exception {
        // bean销毁时关闭客户端
        service.shutdownExecutor();
    }
}
